package parser;

import entity.cwe.DetectionMethod;
import entity.cwe.Weakness;
import org.dom4j.Document;
import org.dom4j.Element;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import util.XmlFileUtil;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * @author xym
 * @date 2020/7/15
 */
public class CweXmlParser {
    private static final Logger log = LoggerFactory.getLogger(CweXmlParser.class);

    /**
     * 对外暴露的解析接口，接收源xml文件输入，返回json文件中包含的攻击模式信息
     *
     * @param fileName xml源文件
     */
    public List<Weakness> execute(String fileName) {
        if (!(fileName.endsWith("xml") || fileName.endsWith("XML"))) {
            log.info("{}格式不正确", fileName);
            throw new RuntimeException("文件格式不正确");
        }
        log.info("开始解析xml文件 {}", fileName);

        Document doc = XmlFileUtil.read(fileName);
        // get the tag <Weakness_Catalog>
        Element root = doc.getRootElement();
        // get the tag <Weaknesses>
        Element weaknesses = root.element("Weaknesses");

        List<Weakness> result = new ArrayList<>();
        // get all children tags <Weakness> in <Weaknesses>
        Iterator<Element> weaknessIterator = weaknesses.elementIterator("Weakness");
        while (weaknessIterator.hasNext()) {
            Element singleWeakness = weaknessIterator.next();
            Weakness weakness = innerWeaknessParsing(singleWeakness);
            if (weakness != null) {
                result.add(weakness);
            }
        }
        log.info("解析完成，共计{}条记录", result.size());
        return result;
    }

    /**
     * parse a single tag <Weakness>
     *
     * @param root tag <Weakness>
     * @return information in tag <Weakness>
     */
    private Weakness innerWeaknessParsing(Element root) {
        if (root == null) {
            return null;
        }
        Weakness weaknessEntity = new Weakness();
        weaknessEntity.setId(Integer.parseInt(root.attributeValue("ID")));
        weaknessEntity.setName(root.attributeValue("Name"));

        Iterator<Element> subNodes = root.elementIterator();
        while (subNodes.hasNext()) {
            Element node = subNodes.next();
            String nodeName = node.getName();
            switch (nodeName) {
                case "Description":
                    weaknessEntity.setDescription(node.getText());
                    break;
                case "Related_Weaknesses":
                    Iterator<Element> relatedWeaknessIterator = node.elementIterator("Related_Weakness");
                    while (relatedWeaknessIterator.hasNext()) {
                        Element relatedAttackPattern = relatedWeaknessIterator.next();
                        String value = relatedAttackPattern.attributeValue("Nature")
                                + "-"
                                + relatedAttackPattern.attributeValue("CWE_ID");
                        weaknessEntity.addRelatedWeakness(value);
                    }
                    break;
                case "Alternate_Terms":
                    Iterator<Element> alternateTermIterator = node.elementIterator("Alternate_Term");
                    while (alternateTermIterator.hasNext()) {
                        Element alternateTerm = alternateTermIterator.next();
                        Element term = alternateTerm.element("Term");
                        weaknessEntity.addAlternateTerm(term.getText());
                    }
                    break;
                case "Modes_Of_Introduction":
                    Iterator<Element> introductionIterator = node.elementIterator("Introduction");
                    while (introductionIterator.hasNext()) {
                        Element introduction = introductionIterator.next();
                        Element phase = introduction.element("Phase");
                        weaknessEntity.addModesOfIntroduction(phase.getText());
                    }
                    break;
                case "Likelihood_Of_Exploit":
                    weaknessEntity.setLikelihoodOfExploit(node.getText());
                    break;
                case "Common_Consequences":
                    Iterator<Element> consequenceIterator = node.elementIterator("Consequence");
                    while (consequenceIterator.hasNext()) {
                        Element introduction = consequenceIterator.next();
                        Element impact = introduction.element("Impact");
                        if (impact != null) {
                            weaknessEntity.addCommonConsequences(impact.getText());
                        }
                    }
                    break;
                case "Potential_Mitigations":
                    Iterator<Element> mitigationIterator = node.elementIterator("Mitigation");
                    while (mitigationIterator.hasNext()) {
                        Element subElement = mitigationIterator.next();
                        Element description = subElement.element("Description");
                        if (description.elementIterator().hasNext()) {
                            Iterator<Element> descriptionIterator = description.elementIterator();
                            StringBuilder sb = new StringBuilder();
                            while (descriptionIterator.hasNext()) {
                                Element value = descriptionIterator.next();
                                sb.append(value.getText()).append(" ");
                            }
                            weaknessEntity.addMitigation(sb.toString().replaceAll("\\n", "").trim());
                        } else {
                            weaknessEntity.addMitigation(description.getText());
                        }
                    }
                    break;
                case "Detection_Methods":
                    Iterator<Element> detectionMethodIterator = node.elementIterator("Detection_Method");
                    while (detectionMethodIterator.hasNext()) {
                        Element detectionMethod = detectionMethodIterator.next();
                        DetectionMethod method = new DetectionMethod();
                        method.setMethodName(detectionMethod.element("Method").getText());

                        Element description = detectionMethod.element("Description");
                        if (description.elementIterator().hasNext()) {
                            Iterator<Element> descriptionIterator = description.elementIterator();
                            StringBuilder sb = new StringBuilder();
                            while (descriptionIterator.hasNext()) {
                                Element value = descriptionIterator.next();
                                sb.append(value.getText()).append(" ");
                            }
                            method.setDescription(sb.toString().replaceAll("\\n", "").trim());
                        } else {
                            method.setDescription(detectionMethod.element("Description").getText());
                        }

                        if (detectionMethod.element("Effectiveness") != null) {
                            method.setEffectiveness(detectionMethod.element("Effectiveness").getText());
                        }

                        weaknessEntity.addDetectionMethods(method);
                    }
                    break;
                default:
                    break;
            }
        }
        return weaknessEntity;
    }
}

